// SPDX-FileCopyrightText: © 2024 Tenstorrent Inc.
//
// SPDX-License-Identifier: Apache-2.0

#include "convert_to_chw_pybind.hpp"

#include <pybind11/pybind11.h>
#include <pybind11/stl.h>

#include "convert_to_chw.hpp"
#include "ttnn-pybind/decorators.hpp"

namespace ttnn::operations::experimental::cnn::detail {

namespace py = pybind11;

void bind_convert_to_chw(py::module& module) {
    using OperationType = decltype(ttnn::experimental::convert_to_chw);

    const auto doc = R"doc(
    Convert a tensor from HWC channel ordering to CHW channel ordering.

    The input tensor is expected to be tiled and height-sharded in L1. The output is a row-major width-sharded tensor.

    The output memory configuration is automatically inferred to create a width-sharded output
    with appropriate shard dimensions based on the input tensor's sharding configuration.
    )doc";

    ttnn::bind_registered_operation(
        module,
        ttnn::experimental::convert_to_chw,
        doc,
        ttnn::pybind_overload_t{
            [](const OperationType& self, const ttnn::Tensor& input, const std::optional<DataType> dtype) {
                return self(input, dtype);
            },
            py::arg("input"),
            py::kw_only(),
            py::arg("dtype") = std::nullopt});
}

}  // namespace ttnn::operations::experimental::cnn::detail
